# agent_loop.py
# 多回合 function-calling 编排：
# user -> deepseek(tool_calls)* -> execute tools -> ... -> deepseek(final)
from __future__ import annotations
import json
from typing import Any, Dict, List

from registry import ToolRegistry, ToolSpec
from executor import Executor
from tools import AddArgs, add_tool, SearchArgs, search_tool, repair_search_args
from deepseek_client import build_tools_payload, deepseek_chat, deepseek_chat_no_tools, MOCK_MODE

def build_registry() -> ToolRegistry:
    reg = ToolRegistry()
    reg.register(ToolSpec(
        name="add",
        params_model=AddArgs,
        func=add_tool,
        desc="加法计算",
        max_retries=0,
    ))
    reg.register(ToolSpec(
        name="search",
        params_model=SearchArgs,
        func=search_tool,
        desc="关键词检索（内置语料）",
        max_retries=2,
        repair=repair_search_args,
    ))
    return reg

def run_agent(user_text: str, max_rounds: int = 5) -> None:
    # 输出当前模式信息，便于调试
    print(f"[DEBUG] Running in {'MOCK' if MOCK_MODE else 'LIVE'} mode")
    
    registry = build_registry()
    executor = Executor(registry)
    tools_payload = build_tools_payload(registry)

    messages: List[Dict[str, Any]] = [
        {"role": "system", "content": (
            "你是一个会合理使用工具的助理。"
            "当用户需求包含多个子任务时，可以在一次或多次 tool_calls 中逐步完成，"
            "直到所有子任务完成为止。若无需工具则直接给出回答。"
            "最终回答阶段请给出整洁的自然语言输出，不要在最终回答里包含任何与工具相关的标记。"
        )},
        {"role": "user", "content": user_text},
    ]

    # 允许模型跨多回合逐步提出工具调用
    for _ in range(max_rounds):
        resp = deepseek_chat(messages, tools_payload)
        msg = resp["choices"][0]["message"]
        tool_calls = msg.get("tool_calls") or []
        content = (msg.get("content") or "").strip() if msg.get("content") else ""

        if not tool_calls:
            # 没有工具调用时，若已给出内容则直接作为最终回答；否则请求一次无工具的总结
            if content:
                print("=== 最终回答 ===")
                print(content)
                return
            final_resp = deepseek_chat_no_tools(messages)
            final_msg = final_resp["choices"][0]["message"]["content"]
            print("=== 最终回答 ===")
            print(final_msg)
            return

        # 执行本轮所有 tool_calls
        for call in tool_calls:
            name = call["function"]["name"]
            args = json.loads(call["function"]["arguments"] or "{}")
            ok, payload, debug = executor.execute(name, args)

            # 调试输出
            print(f"--- ToolCall {name} ---")
            print(debug)
            print("[result]" if ok else "[failed]", payload, "\n")

            # 把工具调用与结果回灌给模型
            messages.append({
                "role": "assistant",
                "tool_calls": [call],
                "content": None
            })
            messages.append({
                "role": "tool",
                "tool_call_id": call["id"],
                "name": name,
                "content": json.dumps({"ok": ok, "data": payload}, ensure_ascii=False)
            })

    # 达到最大轮次仍未给出最终回答，做一次总结后返回
    final_resp = deepseek_chat_no_tools(messages)
    final_msg = final_resp["choices"][0]["message"]["content"]
    print("=== 最终回答 ===")
    print(final_msg)
